#202208030130-大数据2201-张志龙
# 导入相关库
import tushare as ts
import numpy as np
import pandas as pd
import talib
import matplotlib.pyplot as plt
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import accuracy_score, classification_report
from sklearn.model_selection import GridSearchCV, train_test_split
import matplotlib as mpl
plt.rcParams['font.sans-serif'] = ['SimHei']
plt.rcParams['axes.unicode_minus'] = False
mpl.rcParams['font.size'] = 12

# 1. 股票数据获取
# 初始化pro接口
pro = ts.pro_api('1c7f85b9026518588c0d0cdac712c2d17344332c9c8cfe6bc83ee75c')

# 获取贵州茅台(600519.SH)2017-2022年的日线数据
df = pro.daily(ts_code="600519.SH", start_date="20170101", end_date="20221231")
df = df.sort_values('trade_date')
df = df.set_index('trade_date')

# 2. 数据预处理和特征工程
# 转换数据类型
df[['open', 'high', 'low', 'close', 'pre_close', 'change', 'pct_chg', 'vol', 'amount']] = \
    df[['open', 'high', 'low', 'close', 'pre_close', 'change', 'pct_chg', 'vol', 'amount']].apply(pd.to_numeric)

# 构造基本特征
df['close-open'] = (df['close'] - df['open']) / df['open']
df['high-low'] = (df['high'] - df['low']) / df['low']
df['price_change'] = df['close'] - df['pre_close']
df['p_change'] = df['pct_chg']  # 使用API提供的涨跌幅，避免重复计算

# 3. 技术指标计算
# 移动平均线
df['MA5'] = df['close'].rolling(5).mean()
df['MA10'] = df['close'].rolling(10).mean()
df['MA20'] = df['close'].rolling(20).mean()

# TA-Lib指标
df['RSI'] = talib.RSI(df['close'], timeperiod=14)
df['MOM'] = talib.MOM(df['close'], timeperiod=10)
df['EMA12'] = talib.EMA(df['close'], timeperiod=12)
df['EMA26'] = talib.EMA(df['close'], timeperiod=26)
df['MACD'], df['MACDsignal'], df['MACDhist'] = talib.MACD(df['close'], fastperiod=12, slowperiod=26, signalperiod=9)

# 布林带
df['upper'], df['middle'], df['lower'] = talib.BBANDS(df['close'], timeperiod=20)

# 删除缺失值
df.dropna(inplace=True)

# 4. 准备特征和目标变量
# 特征选择
features = ['close', 'vol', 'close-open', 'high-low', 'MA5', 'MA10', 'MA20',
            'RSI', 'MOM', 'EMA12', 'EMA26', 'MACD', 'MACDsignal', 'MACDhist',
            'upper', 'middle', 'lower']

X = df[features]
y = np.where(df['p_change'].shift(-1) > 0, 1, -1)  # 下一天涨为1，跌为-1

# 5. 划分训练集和测试集
# 按时间划分，前80%为训练集，后20%为测试集
split_idx = int(len(X) * 0.8)
X_train, X_test = X.iloc[:split_idx], X.iloc[split_idx:]
y_train, y_test = y[:split_idx], y[split_idx:]

# 6. 构建随机森林模型
# 初始模型
base_model = RandomForestClassifier(n_estimators=100, max_depth=5,
                                   min_samples_leaf=10, random_state=42)
base_model.fit(X_train, y_train)

# 网格搜索优化参数
param_grid = {
    'n_estimators': [50, 100, 150],
    'max_depth': [3, 5, 7],
    'min_samples_leaf': [5, 10, 20]
}

grid_search = GridSearchCV(RandomForestClassifier(random_state=42),
                          param_grid, cv=5, scoring='accuracy', n_jobs=-1)
grid_search.fit(X_train, y_train)

# 最佳模型
best_model = grid_search.best_estimator_
print(f"最佳参数: {grid_search.best_params_}")
print(f"训练集准确率: {accuracy_score(y_train, best_model.predict(X_train)):.4f}")
print(f"测试集准确率: {accuracy_score(y_test, best_model.predict(X_test)):.4f}")

# 7. 模型评价
print("\n分类报告:")
print(classification_report(y_test, best_model.predict(X_test)))

# 在测试数据上添加预测结果
test_df = df.iloc[split_idx:].copy()
test_df['prediction'] = best_model.predict(X_test)

# 计算每天的收益率
test_df['daily_return'] = test_df['pct_chg'] / 100  # 转换为小数形式

# 计算基准收益率 (买入持有策略)
test_df['benchmark'] = (1 + test_df['daily_return']).cumprod()

# 计算策略收益率 (根据模型预测买卖)
test_df['strategy_return'] = (1 + test_df['prediction'].shift(1) * test_df['daily_return']).cumprod()

plt.figure(figsize=(14, 7))
plt.plot(test_df.index, test_df['benchmark'], label='买入持有策略', linewidth=2)
plt.plot(test_df.index, test_df['strategy_return'], label='模型预测策略', linewidth=2)

# 设置图表属性
plt.title('贵州茅台(600519.SH)投资策略收益对比 (2017-2022)', fontsize=16, pad=20)
plt.xlabel('日期', fontsize=14)
plt.ylabel('累计收益率', fontsize=14)
plt.legend(fontsize=12, loc='upper left')
plt.grid(True, linestyle='--', alpha=0.7)

# 优化日期显示
plt.gca().xaxis.set_major_locator(plt.MaxNLocator(10))  # 限制x轴刻度数量
plt.xticks(rotation=45, ha='right')  # 旋转45度并右对齐

# 添加收益标注
final_bench = test_df['benchmark'].iloc[-1] - 1
final_strat = test_df['strategy_return'].iloc[-1] - 1
plt.annotate(f'买入持有: {final_bench:.1%}',
            xy=(test_df.index[-1], test_df['benchmark'].iloc[-1]),
            xytext=(10, 10), textcoords='offset points')
plt.annotate(f'模型策略: {final_strat:.1%}',
            xy=(test_df.index[-1], test_df['strategy_return'].iloc[-1]),
            xytext=(10, -20), textcoords='offset points')

plt.tight_layout()  
plt.savefig('stock_strategy_comparison.png', dpi=300, bbox_inches='tight')  
plt.show()